#include <cstring>
#include <functional>

#include "tinyndarray.h"

namespace tinyndarray {

// 点积运算（优化版本）
NdArray NdArray::dot(const NdArray& other) const {
    if (dtype() != other.dtype()) {
        throw std::runtime_error("Data type mismatch for dot product");
    }

    if (ndim() != 1 || other.ndim() != 1 || size() != other.size()) {
        throw std::runtime_error(
                "Only 1D arrays of same size supported for dot product");
    }

    NdArray result(dtype(), {1});
    const int n = static_cast<int>(size());
    const void* this_ptr = data_ptr();
    const void* other_ptr = other.data_ptr();
    void* result_ptr = result.data_ptr();

    switch (dtype().tag) {
        case FP32: {
            const float* this_data = static_cast<const float*>(this_ptr);
            const float* other_data = static_cast<const float*>(other_ptr);
            float sum = 0.0f;
            for (int i = 0; i < n; ++i) {
                sum += this_data[i] * other_data[i];
            }
            *static_cast<float*>(result_ptr) = sum;
            break;
        }
        case INT32: {
            const int* this_data = static_cast<const int*>(this_ptr);
            const int* other_data = static_cast<const int*>(other_ptr);
            int sum = 0;
            for (int i = 0; i < n; ++i) {
                sum += this_data[i] * other_data[i];
            }
            *static_cast<int*>(result_ptr) = sum;
            break;
        }
        case FP64: {
            const double* this_data = static_cast<const double*>(this_ptr);
            const double* other_data = static_cast<const double*>(other_ptr);
            double sum = 0.0;
            for (int i = 0; i < n; ++i) {
                sum += this_data[i] * other_data[i];
            }
            *static_cast<double*>(result_ptr) = sum;
            break;
        }
        default:
            throw std::runtime_error("Unsupported data type for dot product");
    }

    return result;
}

}  // namespace tinyndarray